import random
import gymnasium as gym
import numpy as np
import collections
from tqdm import tqdm
import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
import os
import json
from datetime import datetime

# 创建保存结果的目录
def create_save_dirs():
    """创建保存模型和结果的目录"""
    base_dir = "dqn_cartpole_results"
    model_dir = os.path.join(base_dir, "models")
    plot_dir = os.path.join(base_dir, "plots")
    
    for directory in [base_dir, model_dir, plot_dir]:
        if not os.path.exists(directory):
            os.makedirs(directory)
    
    return base_dir, model_dir, plot_dir

# 创建目录
BASE_DIR, MODEL_DIR, PLOT_DIR = create_save_dirs()

class ReplayBuffer:
    ''' 经验回放池 '''
    def __init__(self, capacity):
        self.buffer = collections.deque(maxlen=capacity)  # 队列,先进先出

    def add(self, state, action, reward, next_state, done):  # 将数据加入buffer
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size):  # 从buffer中采样数据,数量为batch_size
        transitions = random.sample(self.buffer, batch_size)
        state, action, reward, next_state, done = zip(*transitions)
        return np.array(state), action, reward, np.array(next_state), done

    def size(self):  # 目前buffer中数据的数量
        return len(self.buffer)

class Qnet(torch.nn.Module):
    ''' 只有一层隐藏层的Q网络 '''
    def __init__(self, state_dim, hidden_dim, action_dim):
        super(Qnet, self).__init__()
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim, action_dim)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        return self.fc2(x)

class DQN:
    ''' DQN算法 '''
    def __init__(self, state_dim, hidden_dim, action_dim, learning_rate, gamma,
                 epsilon, target_update, device):
        self.action_dim = action_dim
        self.q_net = Qnet(state_dim, hidden_dim, self.action_dim).to(device)
        self.target_q_net = Qnet(state_dim, hidden_dim, self.action_dim).to(device)
        self.optimizer = torch.optim.Adam(self.q_net.parameters(), lr=learning_rate)
        self.gamma = gamma
        self.epsilon = epsilon
        self.target_update = target_update
        self.count = 0
        self.device = device

    def take_action(self, state):
        if np.random.random() < self.epsilon:
            action = np.random.randint(self.action_dim)
        else:
            state = torch.tensor(np.array([state]), dtype=torch.float).to(self.device)
            action = self.q_net(state).argmax().item()
        return action

    def update(self, transition_dict):
        states = torch.tensor(transition_dict['states'], dtype=torch.float).to(self.device)
        actions = torch.tensor(transition_dict['actions']).view(-1, 1).to(self.device)
        rewards = torch.tensor(transition_dict['rewards'], dtype=torch.float).view(-1, 1).to(self.device)
        next_states = torch.tensor(transition_dict['next_states'], dtype=torch.float).to(self.device)
        dones = torch.tensor(transition_dict['dones'], dtype=torch.float).view(-1, 1).to(self.device)

        q_values = self.q_net(states).gather(1, actions)
        max_next_q_values = self.target_q_net(next_states).max(1)[0].view(-1, 1)
        q_targets = rewards + self.gamma * max_next_q_values * (1 - dones)
        dqn_loss = torch.mean(F.mse_loss(q_values, q_targets))
        
        self.optimizer.zero_grad()
        dqn_loss.backward()
        self.optimizer.step()

        if self.count % self.target_update == 0:
            self.target_q_net.load_state_dict(self.q_net.state_dict())
        self.count += 1

    def save_model(self, filepath):
        """保存模型权重"""
        torch.save({
            'q_net_state_dict': self.q_net.state_dict(),
            'target_q_net_state_dict': self.target_q_net.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict()
        }, filepath)

    def load_model(self, filepath):
        """加载模型权重"""
        checkpoint = torch.load(filepath)
        self.q_net.load_state_dict(checkpoint['q_net_state_dict'])
        self.target_q_net.load_state_dict(checkpoint['target_q_net_state_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

def save_training_plots(return_list, mv_return, episode, prefix=""):
    """保存训练曲线图"""
    episodes_list = list(range(len(return_list)))
    
    plt.figure(figsize=(12, 6))
    
    plt.subplot(1, 2, 1)
    plt.plot(episodes_list, return_list)
    plt.xlabel('Episodes')
    plt.ylabel('Returns')
    plt.title('DQN on CartPole-v1')
    
    plt.subplot(1, 2, 2)
    # 修正：使用与移动平均数据长度匹配的x轴数据
    mv_episodes = list(range(len(mv_return)))
    plt.plot(mv_episodes, mv_return)
    plt.xlabel('Episodes')
    plt.ylabel('Moving Average Returns')
    plt.title('Moving Average (window=9)')
    
    plt.tight_layout()
    
    # 保存图片
    filename = f"{prefix}training_progress_episode_{episode}.png"
    filepath = os.path.join(PLOT_DIR, filename)
    plt.savefig(filepath, dpi=300, bbox_inches='tight')
    plt.close()
    
    return filepath

def save_training_info(hyperparams, return_list, mv_return):
    """保存训练信息和结果"""
    info = {
        'timestamp': datetime.now().strftime("%Y-%m-%d_%H-%M-%S"),
        'hyperparameters': hyperparams,
        'final_results': {
            'last_10_episodes_mean_return': np.mean(return_list[-10:]),
            'last_10_episodes_mean_mv_return': np.mean(mv_return[-10:]),
            'max_return': np.max(return_list),
            'min_return': np.min(return_list),
            'total_episodes': len(return_list)
        }
    }
    
    info_file = os.path.join(BASE_DIR, "training_info.json")
    with open(info_file, 'w') as f:
        json.dump(info, f, indent=4)
    
    return info_file

# 超参数设置
lr = 2e-3
num_episodes = 500
hidden_dim = 128
gamma = 0.98
epsilon = 0.01
target_update = 10
buffer_size = 10000
minimal_size = 500
batch_size = 64
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

# 存储超参数
hyperparams = {
    'learning_rate': lr,
    'num_episodes': num_episodes,
    'hidden_dim': hidden_dim,
    'gamma': gamma,
    'epsilon': epsilon,
    'target_update': target_update,
    'buffer_size': buffer_size,
    'minimal_size': minimal_size,
    'batch_size': batch_size,
    'device': str(device)
}

env_name = 'CartPole-v1'
env = gym.make(env_name)
random.seed(0)
np.random.seed(0)
env.reset(seed=0)
torch.manual_seed(0)

replay_buffer = ReplayBuffer(buffer_size)
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.n
agent = DQN(state_dim, hidden_dim, action_dim, lr, gamma, epsilon, target_update, device)

return_list = []
print(f"开始训练DQN算法在{env_name}环境上...")
print(f"结果将保存在: {BASE_DIR}")

for i in range(10):
    with tqdm(total=int(num_episodes / 10), desc='Iteration %d' % i) as pbar:
        for i_episode in range(int(num_episodes / 10)):
            episode_return = 0
            state, info = env.reset()
            done = False
            
            while not done:
                action = agent.take_action(state)
                next_state, reward, terminated, truncated, _ = env.step(action)
                
                if terminated or truncated:
                    done = True
                
                replay_buffer.add(state, action, reward, next_state, done)
                state = next_state
                episode_return += reward
                
                if replay_buffer.size() > minimal_size:
                    b_s, b_a, b_r, b_ns, b_d = replay_buffer.sample(batch_size)
                    transition_dict = {
                        'states': b_s,
                        'actions': b_a,
                        'next_states': b_ns,
                        'rewards': b_r,
                        'dones': b_d
                    }
                    agent.update(transition_dict)
            
            return_list.append(episode_return)
            
            if (i_episode + 1) % 10 == 0:
                pbar.set_postfix({
                    'episode': '%d' % (num_episodes / 10 * i + i_episode + 1),
                    'return': '%.3f' % np.mean(return_list[-10:])
                })
            
            pbar.update(1)
        
        # 每完成一个迭代保存一次模型和图表
        mv_return = np.convolve(return_list, np.ones(9)/9, mode='valid')
        if len(mv_return) > 0:
            agent.save_model(os.path.join(MODEL_DIR, f"dqn_model_iteration_{i}.pth"))
            plot_path = save_training_plots(return_list, mv_return, (i+1) * int(num_episodes / 10), f"iter_{i}_")

# 训练完成后保存最终结果
mv_return = np.convolve(return_list, np.ones(9)/9, mode='valid')

# 保存最终模型
final_model_path = os.path.join(MODEL_DIR, "dqn_model_final.pth")
agent.save_model(final_model_path)

# 保存最终训练曲线
final_plot_path = save_training_plots(return_list, mv_return, num_episodes, "final_")

# 保存训练信息
info_path = save_training_info(hyperparams, return_list, mv_return)

print(f"\n训练完成！")
print(f"最终模型已保存: {final_model_path}")
print(f"训练曲线图已保存: {final_plot_path}")
print(f"训练信息已保存: {info_path}")
print(f"最后10个回合的平均回报: {np.mean(return_list[-10:]):.3f}")

# 显示最终图表
episodes_list = list(range(len(return_list)))
plt.figure(figsize=(12, 5))

plt.subplot(1, 2, 1)
plt.plot(episodes_list, return_list)
plt.xlabel('Episodes')
plt.ylabel('Returns')
plt.title('DQN on CartPole-v1')

plt.subplot(1, 2, 2)
# 修正：使用与移动平均数据长度匹配的x轴数据
mv_episodes = list(range(len(mv_return)))
plt.plot(mv_episodes, mv_return)
plt.xlabel('Episodes')
plt.ylabel('Moving Average Returns')
plt.title('Moving Average (window=9)')

plt.tight_layout()
plt.savefig(os.path.join(PLOT_DIR, "final_training_results.png"), dpi=300, bbox_inches='tight')
plt.show()